# -*- coding:utf-8-*-
# 指定文件编码为 UTF-8，支持中文字符

# 导入标准库
import json  # 用于处理 JSON 数据
import os  # 提供操作系统相关功能（文件路径等）
import shutil

# 导入 PyTorch 深度学习框架
import torch
import sys  # 提供系统相关功能（如修改系统路径）

# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# 获取 core 文件所在的目录的绝对路径（上一级目录）
rag_qa_path = os.path.dirname(current_dir)
# 获取项目根目录的绝对路径（上两级目录）
project_root = os.path.dirname(rag_qa_path)
# 将项目根目录添加到系统路径，确保可以导入项目中的模块
sys.path.insert(0, project_root)

# 导入自定义日志模块
from base.config import get_config
from base.logger import get_logger

# 导入 numpy 科学计算库
import numpy as np
# 导入 Transformers 库（Hugging Face 的 NLP 库）
from transformers import BertTokenizer, BertForSequenceClassification  # BERT 分词器和分类模型
from transformers import Trainer, TrainingArguments  # 训练工具和参数配置
# 导入 sklearn 的 train_test_split 用于数据分割
from sklearn.model_selection import train_test_split
# 导入 sklearn 的评估指标
from sklearn.metrics import classification_report, confusion_matrix

config = get_config()
logger = get_logger()

class QueryClassifier:
    def __init__(self, model_path="bert_query_classifier"):
        """初始化查询分类器"""
        # 初始化模型保存路径
        self.model_path = model_path
        # 加载 BERT 分词器（中文预训练模型）
        # 构建预训练模型的完整路径
        bert_path = os.path.join(rag_qa_path, 'models', 'bert-base-chinese')
        # 加载预训练的分词器
        self.tokenizer = BertTokenizer.from_pretrained(bert_path)
        # 初始化模型变量（稍后加载）
        self.model = None
        # 确定运行设备（优先使用 GPU，否则使用 CPU）
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # 记录设备信息到日志
        logger.info(f"使用设备: {self.device}")
        # 定义标签映射（文本标签到数字的映射）
        self.label_map = {"通用知识": 0, "专业咨询": 1}
        # 加载模型（如果存在）或初始化新模型
        self.load_model()

    def load_model(self):
        """加载或初始化 BERT 模型"""
        # 检查模型路径是否存在
        if os.path.exists(self.model_path):
            # 加载预训练模型
            self.model = BertForSequenceClassification.from_pretrained(self.model_path)
            # 将模型移动到指定设备（GPU 或 CPU）
            self.model.to(self.device)
            # 记录加载成功的日志
            logger.info(f"加载模型: {self.model_path}")
        else:
            # 初始化新模型（使用预训练权重）
            # 构建预训练模型的完整路径
            bert_path = os.path.join(rag_qa_path, 'models', 'bert-base-chinese')
            self.model = BertForSequenceClassification.from_pretrained(
                bert_path,
                num_labels=2  # 二分类任务
            )
            # 将模型移动到指定设备
            self.model.to(self.device)
            # 记录初始化模型的日志
            logger.info("初始化新 BERT 模型")

    def save_model(self):
        """保存模型和分词器到指定路径"""
        self.save_model_with_cleanup()


        self.model.save_pretrained(self.model_path)
        self.tokenizer.save_pretrained(self.model_path)
        logger.info(f"模型保存至: {self.model_path}")



    def save_model_with_cleanup(self):
        """清理并重新保存模型"""
        try:
            # 1. 先尝试删除已存在的模型目录
            if os.path.exists(self.model_path):
                shutil.rmtree(self.model_path, ignore_errors=True)

            # 2. 重新创建目录
            os.makedirs(self.model_path, exist_ok=True)

            # 3. 保存模型和分词器
            self.model.save_pretrained(self.model_path)
            self.tokenizer.save_pretrained(self.model_path)

            logger.info(f"模型成功保存至: {self.model_path}")

        except PermissionError as e:
            logger.error(f"权限不足，无法保存模型: {e}")
            # 尝试保存到用户目录
            user_model_path = os.path.expanduser(f"~/models/{os.path.basename(self.model_path)}")
            os.makedirs(user_model_path, exist_ok=True)
            self.model.save_pretrained(user_model_path)
            self.tokenizer.save_pretrained(user_model_path)
            logger.info(f"模型已保存至用户目录: {user_model_path}")

        except Exception as e:
            logger.error(f"保存模型时发生错误: {e}")
            raise

    def train_model(self, data_file="model_generic_5000.json"):
        """训练 BERT 分类模型"""
        # 检查数据集文件是否存在
        if not os.path.exists(data_file):
            logger.error(f"数据集文件 {data_file} 不存在")
            raise FileNotFoundError(f"数据集文件 {data_file} 不存在")

        # 加载数据集
        with open(data_file, "r", encoding="utf-8") as f:
            # 读取 JSON 文件中的每一行并解析为字典
            data = [json.loads(line) for line in f.readlines()]

        # 提取查询文本和标签
        texts = [item["query"] for item in data]
        labels = [item["label"] for item in data]

        # 划分训练集和验证集（80%训练，20%验证）
        train_texts, val_texts, train_labels, val_labels = train_test_split(
            texts, labels, test_size=0.2, random_state=42  # 固定随机种子确保可复现性
        )

        # 预处理训练数据（分词、填充等）
        train_encodings, train_labels = self.preprocess_data(train_texts, train_labels)
        # 预处理验证数据
        val_encodings, val_labels = self.preprocess_data(val_texts, val_labels)

        # 创建 PyTorch 数据集对象
        train_dataset = self.create_dataset(train_encodings, train_labels)
        val_dataset = self.create_dataset(val_encodings, val_labels)

        # 配置训练参数
        training_args = TrainingArguments(
            output_dir="./bert_results",  # 输出目录（模型和日志）
            num_train_epochs=3,  # 训练轮次
            per_device_train_batch_size=8,  # 训练批次大小
            per_device_eval_batch_size=8,  # 评估批次大小
            warmup_steps=20,  # 学习率预热步数
            weight_decay=0.01,  # 权重衰减（正则化）
            logging_dir="./bert_logs",  # 日志目录
            logging_steps=10,  # 日志记录频率
            eval_strategy="epoch",  # 每轮结束后评估
            save_strategy="epoch",  # 每轮结束后保存模型
            load_best_model_at_end=True,  # 训练结束时加载最佳模型
            save_total_limit=1,  # 最多保存一个检查点
            metric_for_best_model="eval_loss",  # 选择最佳模型的指标（验证损失）
            fp16=False,  # 禁用混合精度训练
        )

        # 初始化 Trainer 对象
        trainer = Trainer(
            model=self.model,  # 要训练的模型
            args=training_args,  # 训练参数
            train_dataset=train_dataset,  # 训练数据集
            eval_dataset=val_dataset,  # 评估数据集
            compute_metrics=self.compute_metrics  # 计算评估指标的函数
        )

        # 开始训练
        logger.info("开始训练BERT模型")
        trainer.train()

        # 保存训练好的模型
        self.save_model()

        # 在验证集上评估模型性能
        self.evaluate_model(val_texts, val_labels)

    def preprocess_data(self, texts, labels):
        """预处理数据为 BERT 输入格式"""
        # 使用分词器处理文本
        encodings = self.tokenizer(
            texts,
            truncation=True,  # 截断超过最大长度的文本
            padding='max_length',  # 填充到最大长度
            max_length=128,  # 最大序列长度
            return_tensors="pt"  # 返回 PyTorch 张量
        )
        # 将文本标签映射为数字
        labels = [self.label_map[label] for label in labels]
        return encodings, labels

    def create_dataset(self, encodings, labels):
        """创建 PyTorch 数据集对象"""

        # 自定义数据集类
        class Dataset(torch.utils.data.Dataset):
            def __init__(self, encodings, labels):
                super().__init__()
                self.encodings = encodings  # 分词后的输入
                self.labels = labels  # 对应的标签

            def __len__(self):
                """返回数据集大小"""
                return len(self.labels)

            def __getitem__(self, idx):
                """获取单个样本"""
                # 提取指定索引的输入特征
                item = {key: val[idx] for key, val in self.encodings.items()}
                # 添加标签
                item["labels"] = torch.tensor(self.labels[idx])
                return item

        return Dataset(encodings, labels)

    def compute_metrics(self, eval_pred):
        """计算评估指标（用于训练过程中的评估）"""
        logits, labels = eval_pred  # 模型输出和真实标签
        predictions = np.argmax(logits, axis=-1)  # 获取预测类别
        accuracy = (predictions == labels).mean()  # 计算准确率
        return {"accuracy": accuracy}  # 返回指标字典

    def evaluate_model(self, texts, labels):
        """评估模型性能（训练后使用）"""
        # 预处理文本数据
        encodings = self.tokenizer(
            texts,
            truncation=True,
            padding="max_length",
            max_length=128,
            return_tensors="pt"
        )
        # 创建数据集
        dataset = self.create_dataset(encodings, labels)

        # 初始化 Trainer 用于预测
        trainer = Trainer(model=self.model)
        # 进行预测
        predictions = trainer.predict(dataset)
        # 获取预测标签（取 logits 最大值的索引）
        pred_labels = np.argmax(predictions.predictions, axis=-1)
        # 真实标签（已经是数字形式）
        true_labels = labels

        # 记录分类报告（精确率、召回率、F1值等）
        logger.info("分类报告:")
        logger.info(classification_report(
            true_labels,
            pred_labels,
            target_names=["通用知识", "专业咨询"]  # 类别名称
        ))

        # 记录混淆矩阵
        logger.info("混淆矩阵:")
        logger.info(confusion_matrix(true_labels, pred_labels))

    """
        需求：根据输入的查询文本，预测其所属的类别（通用知识或者专业咨询）
            “AI学费一年是多少钱？”->“专业咨询”
        实现思路：
            1：加载模型，并检查模型的状态
                1.1：验证模型是否已经加载，未加载则记录错误并返回默认类别，0-》通用知识，让大模型处理query
            2：输入数据处理
                2.1：对查询语句进行分词和编码（截断/填充到128长度）
                2.2：将编码数据移动到模型所在的设备（CPU/GPU）
            3：执行预测
                3.1：在无梯度模式下执行推理
                3.2：获取模型输出并解析预测结果（取 logits 最大值的索引）
            4：结果映射
                4.1：将模型输出的类别索引映射为类别名称（通用知识/专业咨询）
    """
    def predict_query(self, query):
        # 1：加载模型，并检查模型状态
        if self.model is None:
            logger.error("模型未加载，请先调用train_model方法进行训练")
            return "通用知识"  # 默认类别
        # 2：输入数据处理
        # 2.1：对查询语句进行分词和编码（截断/填充到128长度）
        encodings = self.tokenizer(
            query,
            truncation=True,
            padding="max_length",
            max_length=128,
            return_tensors="pt"
        )
        # 2.2：将编码数据移动到模型所在的设备（CPU/GPU）
        encoding = {k: v.to(self.device) for k, v in encodings.items()}
        # 3：不计算梯度，进行预测
        with torch.no_grad():
            # 获取模型输出（logits）
            outputs = self.model(**encoding)
            # 获取预测结果
            prediction = torch.argmax(outputs.logits, dim=1).item()
        # 4：结果映射
        # 4.1：将模型输出的类别索引映射为类别名称（通用知识/专业咨询）
        return "专业咨询" if prediction ==1 else "通用知识"

if __name__ == '__main__':
    # 创建查询分类器实例
    query_classify = QueryClassifier()
    # 指定训练数据文件路径
    # data_file = 'F://code//integrated_qa_system//rag_qa//data//classify_data//model_generic_5000.json'
    # # 训练模型
    # query_classify.train_model(data_file)
    # 示例，使用训练好的模型进行预测
    result1 = query_classify.predict_query("AI的课程大纲是什么？")
    print(result1)

    result2 = query_classify.predict_query("5*9等于多少？")
    print(result2)
